package org.wheel.plugins.toolkit.util;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Function;

/**
 * 线程安全的LinkedHashMap实现
 * 线程安全且保持插入顺序的Map实现
 * 使用 ConcurrentHashMap + 维护插入顺序
 * 不允许存储 null 值
 * @param <K> 键类型
 * @param <V> 值类型
 */
public class ConcurrentLinkedHashMap<K, V> implements Map<K, V> {
    // 存储实际数据
    private final ConcurrentHashMap<K, V> map = new ConcurrentHashMap<>();
    // 维护键的插入顺序
    private final ConcurrentLinkedQueue<K> order = new ConcurrentLinkedQueue<>();
    // 用于putIfAbsent等操作的同步锁
    private final Object putLock = new Object();

    @Override
    public int size() {
        return map.size();
    }

    @Override
    public boolean isEmpty() {
        return map.isEmpty();
    }

    @Override
    public boolean containsKey(Object key) {
        return map.containsKey(key);
    }

    @Override
    public boolean containsValue(Object value) {
        return map.containsValue(value);
    }

    @Override
    public V get(Object key) {
        return map.get(key);
    }

    @Override
    public V put(K key, V value) {
        synchronized (putLock) {
            if (map.putIfAbsent(key, value) == null) {
                order.add(key);
                return null;
            }
            return map.put(key, value);
        }
    }

    @Override
    public V remove(Object key) {
        synchronized (putLock) {
            V removed = map.remove(key);
            if (removed != null) {
                order.remove(key);
            }
            return removed;
        }
    }

    @Override
    public void putAll(Map<? extends K, ? extends V> m) {
        synchronized (putLock) {
            for (Entry<? extends K, ? extends V> entry : m.entrySet()) {
                put(entry.getKey(), entry.getValue());
            }
        }
    }

    @Override
    public void clear() {
        synchronized (putLock) {
            map.clear();
            order.clear();
        }
    }

    @Override
    public Set<K> keySet() {
        return new LinkedHashSet<>(order);
    }

    @Override
    public Collection<V> values() {
        List<V> values = new ArrayList<>(order.size());
        for (K key : order) {
            values.add(map.get(key));
        }
        return Collections.unmodifiableList(values);
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        Set<Entry<K, V>> entries = new LinkedHashSet<>();
        for (K key : order) {
            entries.add(new AbstractMap.SimpleEntry<>(key, map.get(key)));
        }
        return Collections.unmodifiableSet(entries);
    }

    @Override
    public V getOrDefault(Object key, V defaultValue) {
        V v;
        return ((v = get(key)) != null) ? v : defaultValue;
    }

    @Override
    public void forEach(BiConsumer<? super K, ? super V> action) {
        Objects.requireNonNull(action);
        for (K key : order) {
            V value = map.get(key);
            action.accept(key, value);
        }
    }

    @Override
    public V putIfAbsent(K key, V value) {
        synchronized (putLock) {
            V v = map.get(key);
            if (v == null) {
                v = map.putIfAbsent(key, value);
                if (v == null) {
                    order.add(key);
                }
            }
            return v;
        }
    }

    @Override
    public boolean remove(Object key, Object value) {
        synchronized (putLock) {
            boolean removed = map.remove(key, value);
            if (removed) {
                order.remove(key);
            }
            return removed;
        }
    }

    @Override
    public boolean replace(K key, V oldValue, V newValue) {
        synchronized (putLock) {
            boolean replaced = map.replace(key, oldValue, newValue);
            if (replaced && !order.contains(key)) {
                order.add(key);
            }
            return replaced;
        }
    }

    @Override
    public V replace(K key, V value) {
        synchronized (putLock) {
            V replaced = map.replace(key, value);
            if (replaced != null && !order.contains(key)) {
                order.add(key);
            }
            return replaced;
        }
    }

    @Override
    public void replaceAll(BiFunction<? super K, ? super V, ? extends V> function) {
        synchronized (putLock) {
            for (K key : order) {
                map.computeIfPresent(key, function);
            }
        }
    }

    @Override
    public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
        Objects.requireNonNull(mappingFunction);
        synchronized (putLock) {
            V v = map.get(key);
            if (v == null) {
                V newValue = mappingFunction.apply(key);
                if (newValue != null) {
                    map.putIfAbsent(key, newValue);
                    order.add(key);
                    return newValue;
                }
            }
            return v;
        }
    }

    @Override
    public V computeIfPresent(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
        synchronized (putLock) {
            V oldValue = map.get(key);
            if (oldValue != null) {
                V newValue = remappingFunction.apply(key, oldValue);
                if (newValue != null) {
                    map.replace(key, newValue);
                    if (!order.contains(key)) {
                        order.add(key);
                    }
                    return newValue;
                } else {
                    map.remove(key);
                    order.remove(key);
                    return null;
                }
            }
            return null;
        }
    }

    @Override
    public V compute(K key, BiFunction<? super K, ? super V, ? extends V> remappingFunction) {
        synchronized (putLock) {
            V oldValue = map.get(key);
            V newValue = remappingFunction.apply(key, oldValue);
            if (newValue == null) {
                // 删除映射
                if (oldValue != null) {
                    map.remove(key);
                    order.remove(key);
                }
                return null;
            } else {
                // 添加或替换
                map.put(key, newValue);
                if (oldValue == null) {
                    order.add(key);
                } else if (!order.contains(key)) {
                    order.add(key);
                }
                return newValue;
            }
        }
    }

    @Override
    public V merge(K key, V value, BiFunction<? super V, ? super V, ? extends V> remappingFunction) {
        synchronized (putLock) {
            V oldValue = map.get(key);
            V newValue = (oldValue == null) ? value :
                    remappingFunction.apply(oldValue, value);
            if (newValue == null) {
                map.remove(key);
                order.remove(key);
            } else {
                map.put(key, newValue);
                if (oldValue == null) {
                    order.add(key);
                }
            }
            return newValue;
        }
    }

    /**
     * 获取按照插入顺序排列的键列表
     */
    public List<K> orderedKeys() {
        return new ArrayList<>(order);
    }

    /**
     * 获取按照插入顺序排列的条目列表
     */
    public List<Entry<K, V>> orderedEntries() {
        List<Entry<K, V>> entries = new ArrayList<>(order.size());
        for (K key : order) {
            entries.add(new AbstractMap.SimpleEntry<>(key, map.get(key)));
        }
        return entries;
    }
}